package top.jolyoulu.corecommon.algorithms.flowcontrol;

import java.time.Instant;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @Author LuZhouJin
 * @Date 2023/6/1
 * 漏桶限流算法
 * 思路：
 * 当前有一个桶，桶中有一个洞会根据一定的速率流水桶中的水，
 * 在需要处理任务时前往桶中添加相当比例的水量
 * 假如桶溢出，代表已经满负荷运作无法接受更多的任务了
 * 假如桶未溢出，代表还有空余位置处理任务
 */
public class LeakyBucket {
    //桶的容量
    private final int capacity;
    //桶中漏水速度(单位水量)，即每秒桶中会漏多少水量
    private final int rate;
    //当前桶的水量，当水量达到capacity表示桶已满
    private final AtomicInteger water;
    //上次计算水量的时间
    //(当前时间 - 上一次时间) * rate = 这次water需要减少多少水量
    private Instant lastLeakTime;

    public LeakyBucket(int capacity, int rate) {
        this.capacity = capacity;
        this.rate = rate;
        this.water = new AtomicInteger(0);
        this.lastLeakTime = Instant.now();
    }

    /**
     * 尝试往桶中添加 1 水量
     * @return true成功，false失败
     */
    public boolean acquire(){
        return acquire(1);
    }

    /**
     * 尝试往桶中添加指定水量
     * @param tokens 水量
     * @return true成功，false失败
     */
    public boolean acquire(int tokens) {
        //更新水量
        leakWater();
        //原子性地增加水量，并获取当前水量
        int currentWaterLevel = water.addAndGet(tokens);
        //判断添加后水量是否超出capacity
        if (currentWaterLevel <= capacity) {
            return true;
        } else {
            //水量超过容量，回滚水量的增加
            water.getAndAdd(-tokens);
            return false;
        }
    }

    private void leakWater() {
        //获取当前时间
        Instant now = Instant.now();
        //当前时间 - 上一次更新时间 = 间隔多少秒
        long seconds = now.getEpochSecond() - lastLeakTime.getEpochSecond();
        //间隔秒 * 漏水速度 = 这段间隔时间桶需要漏掉多少水量
        double leakAmount = seconds * rate;
        if (leakAmount > 0) {
            //获取桶当前的水量
            int currentWaterLevel = water.get();
            //当前水量 - 这段间隔时间桶需要漏掉的水量，让桶的水位下降
            //利用Math.max方法使水位不会出现负数
            int newWater = Math.max(0, currentWaterLevel - (int) leakAmount);
            //将新计算的水量设置到当前桶中，设置当前最后计算时间
            water.set(newWater);
            lastLeakTime = now;
        }
    }


    public static void main(String[] args) throws InterruptedException {
        LeakyBucket bucket = new LeakyBucket(1, 1);
        List<Thread> threads = new ArrayList<>();
        for (int i = 0; i < 2; i++) {
            threads.add(new Thread(() -> {
                while (true){
                    String name = Thread.currentThread().getName();
                    System.out.println(name+"=> 尝试往桶里加水，即申请令牌");
                    if (bucket.acquire(1)) {
                        System.out.println(name+"=> 申请成功");
                    }else {
                        System.out.println(name+"=> 申请失败");
                    }
                    try {
                        TimeUnit.MILLISECONDS.sleep(500); //休眠1秒
                    } catch (InterruptedException e) {
                        Thread.currentThread().interrupt();
                    }
                }
            },"Thread-"+i));
        }
        threads.forEach(Thread::start);
        TimeUnit.SECONDS.sleep(60);
    }
}
